import numpy as np

class PIDController:
    """
    A PID (Proportional-Integral-Derivative) controller implementation.
    Used to compute control outputs for systems based on proportional, integral, and derivative terms.

    Parameters:
    kp -- Proportional gain.
    ki -- Integral gain.
    kd -- Derivative gain.

    Methods:
    compute(error, dt) -- Compute the control output based on the current error and time step.
    """
    def __init__(self, kp, ki, kd):
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.prev_error = 0
        self.integral = 0

    def compute(self, error, dt):
        """
        Compute the PID control output.

        Parameters:
        error -- Current error signal.
        dt -- Time step.

        Returns:
        PID control output.
        """
        proportional = self.kp * error
        self.integral += error * dt
        integral = self.ki * self.integral
        derivative = self.kd * (error - self.prev_error) / dt
        self.prev_error = error
        return proportional + integral + derivative


# Instantiate PID controllers for steering and speed control
steering_pid = PIDController(kp=1.0, ki=0.0, kd=0.1)
speed_pid = PIDController(kp=1.0, ki=0.0, kd=0.0)
dt = 0.01  # Time step for PID computations


def calculate_steering_angle(vehicle_pos, vehicle_orientation, target_pos):
    """
    Calculate the steering angle required to steer toward a target position.

    Parameters:
    vehicle_pos -- Current position of the vehicle [x, y].
    vehicle_orientation -- Current orientation of the vehicle (in radians).
    target_pos -- Target position [x, y].

    Returns:
    Angle difference between the target direction and vehicle orientation (in radians).
    """
    dx = target_pos[0] - vehicle_pos[0]
    dy = target_pos[1] - vehicle_pos[1]
    target_angle = np.arctan2(dy, dx)
    angle_diff = target_angle - vehicle_orientation
    angle_diff = (angle_diff + np.pi) % (2 * np.pi) - np.pi  # Normalize to [-pi, pi]
    return angle_diff


def select_closest_target_with_lookahead(vehicle_pos, vehicle_orientation, path, lookahead_distance):
    """
    Select the closest target point within a specified lookahead distance.

    Parameters:
    vehicle_pos -- Current position of the vehicle [x, y].
    vehicle_orientation -- Current orientation of the vehicle (in radians).
    path -- List of points representing the path, where each point is [x, y, theta, kappa].
    lookahead_distance -- Minimum lookahead distance for selecting the target point.

    Returns:
    Closest target point [x, y, theta, kappa] or the final point in the path if no suitable point is found.
    """
    closest_point = None
    min_distance = float('inf')

    for point in path:
        distance = np.linalg.norm(vehicle_pos - point[:2])

        dx = point[0] - vehicle_pos[0]
        dy = point[1] - vehicle_pos[1]
        target_angle = np.arctan2(dy, dx)
        angle_diff = target_angle - vehicle_orientation
        angle_diff = (angle_diff + np.pi) % (2 * np.pi) - np.pi  # Normalize to [-pi, pi]

        if abs(angle_diff) <= np.pi / 2 and lookahead_distance <= distance < min_distance:
            min_distance = distance
            closest_point = point

    # If no suitable point is found, return the last point in the path
    if np.linalg.norm(vehicle_pos - np.array([path[-1][0], path[-1][1]])) < 3:
        print("Returning the final path point")
        return path[-1]

    if closest_point is None and path:
        print("Returning the final path point")
        closest_point = path[-1]

    return closest_point


def pid_controller_incremental(vehicle_pos, vehicle_orientation, current_steering_angle, current_speed, path, 
                                max_speed=4.0, alpha=0.8, lookahead_distance=3.0, print_log = True):
    """
    Incremental PID controller for dynamic path following.

    Parameters:
    vehicle_pos -- Current position of the vehicle [x, y].
    vehicle_orientation -- Current orientation of the vehicle (in radians).
    current_steering_angle -- Current steering angle of the vehicle (in radians).
    current_speed -- Current speed of the vehicle.
    path -- List of points representing the path, where each point is [x, y, theta, kappa].
    max_speed -- Maximum allowed speed (default: 4.0).
    alpha -- Weight factor for adjusting speed based on angle error (default: 0.8).
    lookahead_distance -- Lookahead distance for selecting the target point (default: 3.0).
    print_log -- Print debug information.

    Returns:
    control_speed -- Computed speed control output.
    control_steering -- Computed steering control output.
    angle_error -- Error between the target angle and vehicle orientation.
    target_speed -- Target speed based on path and geometry.
    """
    if not path:
        print('Path is empty.')
        return 0, 0

    target_point = select_closest_target_with_lookahead(vehicle_pos, vehicle_orientation, path, lookahead_distance)
    if target_point is None:
        print('Target point not found.')
        return 0, 0

    x_target, y_target, yaw_target, k_target = target_point

    dx = x_target - vehicle_pos[0]
    dy = y_target - vehicle_pos[1]
    geometric_target_angle = np.arctan2(dy, dx)
    target_angle = geometric_target_angle

    angle_error = target_angle - vehicle_orientation
    angle_error = (angle_error + np.pi) % (2 * np.pi) - np.pi  # Normalize to [-pi, pi]

    delta_steering = steering_pid.compute(angle_error, dt)
    control_steering = current_steering_angle + delta_steering
    control_steering = np.clip(control_steering, -np.radians(49), np.radians(49))

    target_speed = max_speed * (1 - alpha * abs(angle_error))
    speed_error = target_speed - current_speed
    control_speed = speed_pid.compute(speed_error, dt)
    control_speed = np.clip(control_speed, 0.4, max_speed)
    
    if print_log:
        print(f'Vehicle: [{vehicle_pos} {vehicle_orientation}], Target point: [{x_target} {y_target}] {yaw_target}')
        print(f'Angle error: {angle_error}, Delta steering: {delta_steering}, Control steering: {control_steering}')
        print(f'Speed error: {speed_error}, Control speed: {control_speed}')
        print(' ')

    return control_speed, control_steering, angle_error, target_speed


def pid_controller(vehicle_pos, vehicle_orientation, current_speed, path, max_speed=4.0, alpha=0.5, lookahead_distance=3.0):
    """
    Geometric steering-based PID controller for path following.

    Parameters:
    vehicle_pos -- Current position of the vehicle [x, y].
    vehicle_orientation -- Current orientation of the vehicle (in radians).
    current_speed -- Current speed of the vehicle.
    path -- List of points representing the path, where each point is [x, y, theta, kappa].
    max_speed -- Maximum allowed speed (default: 4.0).
    alpha -- Weight factor for adjusting speed based on angle error (default: 0.5).
    lookahead_distance -- Lookahead distance for selecting the target point (default: 3.0).

    Returns:
    control_speed -- Computed speed control output.
    control_steering -- Computed steering control output.
    target_angle -- Target angle for the vehicle's orientation.
    target_speed -- Target speed based on path and geometry.
    """
    if not path:
        print('Path is empty.')
        return 0, 0

    target_point = select_closest_target_with_lookahead(vehicle_pos, vehicle_orientation, path, lookahead_distance)
    if target_point is None:
        print('Target point not found.')
        return 0, 0

    x_target, y_target, yaw_target, k_target = target_point

    dx = x_target - vehicle_pos[0]
    dy = y_target - vehicle_pos[1]
    target_angle = np.arctan2(dy, dx)
    angle_error = target_angle - vehicle_orientation
    angle_error = (angle_error + np.pi) % (2 * np.pi) - np.pi  # Normalize to [-pi, pi]

    pid_steering = steering_pid.compute(angle_error, dt)
    control_steering = np.clip(pid_steering, -np.radians(49), np.radians(49))

    target_speed = max_speed * (1 - alpha * abs(angle_error))
    speed_error = target_speed - current_speed
    control_speed = speed_pid.compute(speed_error, dt)
    control_speed = np.clip(control_speed, 0.0, max_speed)

    print(f'Target angle: {target_angle}, Target speed: {target_speed}')
    print(f'Angle error: {angle_error}, Control steering: {control_steering}')
    print(f'Speed error: {speed_error}, Control speed: {control_speed}')
    print(' ')

    return control_speed, control_steering, target_angle, target_speed
